import cv2
import numpy as np
import random
from typing import Tuple, Optional

class DataAugmentation:
    """
    数据增强类，包含镜像、旋转、亮度、对比度、随机噪声等增强方法
    专门为边缘检测任务设计，确保图像和标签同步变换
    """
    
    def __init__(self, 
                 enable_flip: bool = True,
                 enable_rotation: bool = True, 
                 enable_brightness: bool = True,
                 enable_contrast: bool = True,
                 enable_noise: bool = True,
                 flip_prob: float = 0.5,
                 rotation_prob: float = 0.3,
                 brightness_prob: float = 0.3,
                 contrast_prob: float = 0.3,
                 noise_prob: float = 0.2,
                 rotation_range: float = 15.0,
                 brightness_range: float = 0.2,
                 contrast_range: float = 0.2,
                 noise_std: float = 10.0):
        """
        初始化数据增强参数
        
        Args:
            enable_*: 是否启用各种增强
            *_prob: 各种增强的应用概率
            rotation_range: 旋转角度范围 (-range, +range)
            brightness_range: 亮度调整范围 (-range, +range)
            contrast_range: 对比度调整范围 (1-range, 1+range)
            noise_std: 高斯噪声标准差
        """
        self.enable_flip = enable_flip
        self.enable_rotation = enable_rotation
        self.enable_brightness = enable_brightness
        self.enable_contrast = enable_contrast
        self.enable_noise = enable_noise
        
        self.flip_prob = flip_prob
        self.rotation_prob = rotation_prob
        self.brightness_prob = brightness_prob
        self.contrast_prob = contrast_prob
        self.noise_prob = noise_prob
        
        self.rotation_range = rotation_range
        self.brightness_range = brightness_range
        self.contrast_range = contrast_range
        self.noise_std = noise_std
    
    def horizontal_flip(self, img: np.ndarray, gt: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
        """水平镜像翻转"""
        img_flipped = cv2.flip(img, 1)
        gt_flipped = cv2.flip(gt, 1)
        
        # 确保灰度图像保持3维形状 (H, W, 1)
        if len(img.shape) == 3 and img.shape[-1] == 1 and len(img_flipped.shape) == 2:
            img_flipped = np.expand_dims(img_flipped, axis=-1)
        
        return img_flipped, gt_flipped
    
    def vertical_flip(self, img: np.ndarray, gt: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
        """垂直镜像翻转"""
        img_flipped = cv2.flip(img, 0)
        gt_flipped = cv2.flip(gt, 0)
        
        # 确保灰度图像保持3维形状 (H, W, 1)
        if len(img.shape) == 3 and img.shape[-1] == 1 and len(img_flipped.shape) == 2:
            img_flipped = np.expand_dims(img_flipped, axis=-1)
        
        return img_flipped, gt_flipped
    
    def rotate(self, img: np.ndarray, gt: np.ndarray, angle: float) -> Tuple[np.ndarray, np.ndarray]:
        """
        旋转图像和标签
        
        Args:
            img: 输入图像
            gt: 对应的边缘标签
            angle: 旋转角度（度）
        """
        h, w = img.shape[:2]
        center = (w // 2, h // 2)
        
        # 计算旋转矩阵
        rotation_matrix = cv2.getRotationMatrix2D(center, angle, 1.0)
        
        # 旋转图像
        img_rotated = cv2.warpAffine(img, rotation_matrix, (w, h), 
                                    flags=cv2.INTER_LINEAR,
                                    borderMode=cv2.BORDER_REFLECT_101)
        
        # 旋转标签（使用最近邻插值避免产生新的灰度值）
        gt_rotated = cv2.warpAffine(gt, rotation_matrix, (w, h),
                                   flags=cv2.INTER_NEAREST,
                                   borderMode=cv2.BORDER_REFLECT_101)
        
        # 确保灰度图像保持3维形状 (H, W, 1)
        if len(img.shape) == 3 and img.shape[-1] == 1 and len(img_rotated.shape) == 2:
            img_rotated = np.expand_dims(img_rotated, axis=-1)
        
        return img_rotated, gt_rotated
    
    def adjust_brightness(self, img: np.ndarray, delta: float) -> np.ndarray:
        """
        调整图像亮度
        
        Args:
            img: 输入图像
            delta: 亮度调整值 (-1.0 到 1.0)
        """
        # 确保数据类型为float32以避免溢出
        img_float = img.astype(np.float32)
        
        # 调整亮度（加法操作）
        img_bright = img_float + delta * 255.0
        
        # 限制到有效范围
        img_bright = np.clip(img_bright, 0, 255)
        
        return img_bright.astype(img.dtype)
    
    def adjust_contrast(self, img: np.ndarray, alpha: float) -> np.ndarray:
        """
        调整图像对比度
        
        Args:
            img: 输入图像
            alpha: 对比度因子 (0.5 到 2.0)
        """
        # 确保数据类型为float32
        img_float = img.astype(np.float32)
        
        # 对比度调整：I_new = alpha * I + (1-alpha) * mean
        mean_value = np.mean(img_float)
        img_contrast = alpha * img_float + (1 - alpha) * mean_value
        
        # 限制到有效范围
        img_contrast = np.clip(img_contrast, 0, 255)
        
        return img_contrast.astype(img.dtype)
    
    def add_gaussian_noise(self, img: np.ndarray, std: float) -> np.ndarray:
        """
        添加高斯噪声
        
        Args:
            img: 输入图像
            std: 噪声标准差
        """
        # 生成高斯噪声
        noise = np.random.normal(0, std, img.shape).astype(np.float32)
        
        # 添加噪声
        img_noisy = img.astype(np.float32) + noise
        
        # 限制到有效范围
        img_noisy = np.clip(img_noisy, 0, 255)
        
        return img_noisy.astype(img.dtype)
    
    def random_augment(self, img: np.ndarray, gt: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
        """
        随机应用数据增强
        
        Args:
            img: 输入图像 (H, W, C)
            gt: 对应的边缘标签 (H, W)
            
        Returns:
            增强后的图像和标签
        """
        aug_img = img.copy()
        aug_gt = gt.copy()
        
        # 1. 随机水平翻转
        if self.enable_flip and random.random() < self.flip_prob:
            flip_type = random.choice(['horizontal', 'vertical', 'both'])
            if flip_type == 'horizontal':
                aug_img, aug_gt = self.horizontal_flip(aug_img, aug_gt)
            elif flip_type == 'vertical':
                aug_img, aug_gt = self.vertical_flip(aug_img, aug_gt)
            else:  # both
                aug_img, aug_gt = self.horizontal_flip(aug_img, aug_gt)
                aug_img, aug_gt = self.vertical_flip(aug_img, aug_gt)
        
        # 2. 随机旋转
        if self.enable_rotation and random.random() < self.rotation_prob:
            angle = random.uniform(-self.rotation_range, self.rotation_range)
            aug_img, aug_gt = self.rotate(aug_img, aug_gt, angle)
        
        # 3. 随机亮度调整（只对图像，不影响标签）
        if self.enable_brightness and random.random() < self.brightness_prob:
            delta = random.uniform(-self.brightness_range, self.brightness_range)
            aug_img = self.adjust_brightness(aug_img, delta)
        
        # 4. 随机对比度调整（只对图像，不影响标签）
        if self.enable_contrast and random.random() < self.contrast_prob:
            alpha = random.uniform(1 - self.contrast_range, 1 + self.contrast_range)
            aug_img = self.adjust_contrast(aug_img, alpha)
        
        # 5. 随机噪声（只对图像，不影响标签）
        if self.enable_noise and random.random() < self.noise_prob:
            noise_std = random.uniform(0, self.noise_std)
            aug_img = self.add_gaussian_noise(aug_img, noise_std)
        
        return aug_img, aug_gt
    
    def __repr__(self):
        """返回增强器的配置信息"""
        config = []
        if self.enable_flip:
            config.append(f"Flip(p={self.flip_prob})")
        if self.enable_rotation:
            config.append(f"Rotation(p={self.rotation_prob}, range=±{self.rotation_range}°)")
        if self.enable_brightness:
            config.append(f"Brightness(p={self.brightness_prob}, range=±{self.brightness_range})")
        if self.enable_contrast:
            config.append(f"Contrast(p={self.contrast_prob}, range=±{self.contrast_range})")
        if self.enable_noise:
            config.append(f"Noise(p={self.noise_prob}, std≤{self.noise_std})")
        
        return f"DataAugmentation({', '.join(config)})"


def create_augmentation_from_args(args) -> Optional[DataAugmentation]:
    """
    根据命令行参数创建数据增强器
    
    Args:
        args: 包含增强参数的argparse对象
        
    Returns:
        DataAugmentation实例或None（如果禁用增强）
    """
    if not getattr(args, 'use_augmentation', False):
        return None
    
    return DataAugmentation(
        enable_flip=getattr(args, 'aug_flip', True),
        enable_rotation=getattr(args, 'aug_rotation', True),
        enable_brightness=getattr(args, 'aug_brightness', True),
        enable_contrast=getattr(args, 'aug_contrast', True),
        enable_noise=getattr(args, 'aug_noise', True),
        flip_prob=getattr(args, 'aug_flip_prob', 0.5),
        rotation_prob=getattr(args, 'aug_rotation_prob', 0.3),
        brightness_prob=getattr(args, 'aug_brightness_prob', 0.3),
        contrast_prob=getattr(args, 'aug_contrast_prob', 0.3),
        noise_prob=getattr(args, 'aug_noise_prob', 0.2),
        rotation_range=getattr(args, 'aug_rotation_range', 15.0),
        brightness_range=getattr(args, 'aug_brightness_range', 0.2),
        contrast_range=getattr(args, 'aug_contrast_range', 0.2),
        noise_std=getattr(args, 'aug_noise_std', 10.0)
    )